package main

import (
	"fmt"
	"math"
)

func main() {
	isBalance := isBalanced(&TreeNode{
		Val: 3,
		Right: &TreeNode{
			Val: 20,
			Left: &TreeNode{
				Val: 15,
			},
			Right: &TreeNode{
				Val: 7,
			},
		},
	})
	fmt.Println(isBalance)
	//res := levelOrder(sortedArrayToBST([]int{-10, -3, 0, 5, 9}))
	//for _, i := range res {
	//	for _, j := range i {
	//		fmt.Print(" ", j)
	//	}
	//	fmt.Println()
	//}
}

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func sortedArrayToBST(nums []int) *TreeNode {
	if len(nums) == 0 {
		return nil
	} else if len(nums) == 1 {
		return &TreeNode{Val: nums[0]}
	}
	mid := len(nums) / 2
	root := &TreeNode{
		Val:   nums[mid],
		Left:  sortedArrayToBST(nums[:mid]),
		Right: sortedArrayToBST(nums[mid+1:]),
	}
	return root
}

func levelOrder(root *TreeNode) (ans [][]int) {
	if root == nil {
		return nil
	}
	var queue [](*TreeNode)
	queue = append(queue, root)
	var level int
	for length := len(queue); length != 0; length = len(queue) {
		ans = append(ans, nil)
		for _, v := range queue {
			ans[level] = append(ans[level], v.Val)
			if v.Left != nil {
				queue = append(queue, v.Left)
			}
			if v.Right != nil {
				queue = append(queue, v.Right)
			}
		}
		queue = queue[length:]
		level++
	}
	return ans
}

func isBalanced(root *TreeNode) bool {
	balance, _ := balanceHelper(root)
	return balance
}
func balanceHelper(root *TreeNode) (bool, int) {
	if root == nil {
		return true, 0
	}
	balanceL, depthL := balanceHelper(root.Left)
	if !balanceL {
		return false, depthL
	}
	balanceR, depthR := balanceHelper(root.Right)
	if !balanceR {
		return false, depthR
	}
	return int(math.Abs(float64(depthR-depthL))) < 2, int(math.Max(float64(depthR), float64(depthL))) + 1
}
